package support.limit.ratelimit.impl;

import cn.hutool.core.util.NumberUtil;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.util.StringUtils;
import support.limit.enums.RateLimitConstantEnums;
import support.limit.exception.RateLimitException;
import support.limit.model.LimitDTO;
import support.limit.ratelimit.AbstractRateLimiter;
import support.limit.utils.KeyUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 令牌桶限流
 *
 * @Description RedisRateLimiterTokenBucketImpl
 * @Author guoyb
 * @Date 2023/12/22 14:23
 */
@Slf4j
public class RedisRateLimiterTokenBucketImpl extends AbstractRateLimiter {

    @Autowired
    private RedisTemplate redisTemplate;

    private DefaultRedisScript redisScript;

    public RedisRateLimiterTokenBucketImpl(DefaultRedisScript redisScript) {
        this.redisScript = redisScript;
    }

    @Override
    public void rateLimit(LimitDTO limitDTO) {

        List<Object> keyList = new ArrayList<>();
        keyList.add(key);
        keyList.add(limit + RateLimitConstantEnums.HASH_TAG);
        keyList.add(tokenBucketStepNum + RateLimitConstantEnums.HASH_TAG);
        keyList.add(tokenBucketTimeInterval + RateLimitConstantEnums.HASH_TAG);
        keyList.add(System.currentTimeMillis() / 1000 + RateLimitConstantEnums.HASH_TAG);
        String toString = redisTemplate.execute(redisScript, keyList, keyList).toString();
        if (RateLimitConstantEnums.REDIS_ERROR.equals(toString)) {
            throw new RateLimitException(message);
        }

    }


    private Map<String, RateLimiter> rateLimiterMap = Maps.newConcurrentMap();

    /**
     * guava实现的令牌桶
     * @param limitEntity
     * @return
     */
    public boolean tryAccess(LimitDTO limitEntity) {
        RateLimiter rateLimiter = getRateLimiter(limitEntity);
        if (rateLimiter == null) {
            return false;
        }
        //添加令牌到桶里
        boolean access = rateLimiter.tryAcquire(1, 2000, TimeUnit.MILLISECONDS);
        return access;
    }

    private RateLimiter getRateLimiter(LimitDTO limitDTO) {
        if (limitDTO == null) {
            return null;
        }
        String key = KeyUtil.compositeKey(limitDTO.getIdentifier(), limitDTO.getKey());
        if (StringUtils.isEmpty(key)) {
            return null;
        }
        RateLimiter rateLimiter = rateLimiterMap.get(key);
        if (rateLimiter == null) {
            long limitNum = limitDTO.getTokenBucketStepNum();
            double permitsPerSecond = NumberUtil.div(limitNum,limitDTO.getTokenBucketTimeInterval());
            RateLimiter newRateLimiter = RateLimiter.create(permitsPerSecond);
            rateLimiter = rateLimiterMap.putIfAbsent(key, newRateLimiter);
            if (rateLimiter == null) {
                rateLimiter = newRateLimiter;
            }
        }
        return rateLimiter;
    }

}
